package kvdbmongo

import (
	"context"
	"io"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"

	"github.com/xiaonanln/goworld/engine/gwlog"
	"github.com/xiaonanln/goworld/engine/kvdb/types"
)

const (
	_DEFAULT_DB_NAME = "goworld"
	_VAL_KEY         = "_"
)

type mongoKVDB struct {
	client     *mongo.Client
	collection *mongo.Collection
}

// OpenMongoKVDB opens mongodb as KVDB engine
func OpenMongoKVDB(url string, dbname string, collectionName string) (kvdbtypes.KVDBEngine, error) {
	gwlog.Debugf("Connecting MongoDB ...")
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	client, err := mongo.Connect(ctx, options.Client().ApplyURI(url))
	if err != nil {
		return nil, err
	}

	if dbname == "" {
		// if db is not specified, use default
		dbname = _DEFAULT_DB_NAME
	}
	db := client.Database(dbname)
	collection := db.Collection(collectionName)
	return &mongoKVDB{
		client:     client,
		collection: collection,
	}, nil
}

func (kvdb *mongoKVDB) Put(key string, val string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	filter := bson.M{"_id": key}
	update := bson.M{"$set": bson.M{_VAL_KEY: val}}
	opts := options.Update().SetUpsert(true)
	_, err := kvdb.collection.UpdateOne(ctx, filter, update, opts)
	return err
}

func (kvdb *mongoKVDB) Get(key string) (val string, err error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	filter := bson.M{"_id": key}
	var doc bson.M
	err = kvdb.collection.FindOne(ctx, filter).Decode(&doc)
	if err != nil {
		if err == mongo.ErrNoDocuments {
			err = nil
		}
		return
	}
	val = doc[_VAL_KEY].(string)
	return
}

type mongoKVIterator struct {
	cursor *mongo.Cursor
	ctx    context.Context
}

func (it *mongoKVIterator) Next() (kvdbtypes.KVItem, error) {
	var doc bson.M
	if it.cursor.Next(it.ctx) {
		err := it.cursor.Decode(&doc)
		if err != nil {
			return kvdbtypes.KVItem{}, err
		}
		return kvdbtypes.KVItem{
			Key: doc["_id"].(string),
			Val: doc[_VAL_KEY].(string),
		}, nil
	}

	if err := it.cursor.Err(); err != nil {
		return kvdbtypes.KVItem{}, err
	}
	return kvdbtypes.KVItem{}, io.EOF
}

func (kvdb *mongoKVDB) Find(beginKey string, endKey string) (kvdbtypes.Iterator, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	filter := bson.M{"_id": bson.M{"$gte": beginKey, "$lt": endKey}}
	cursor, err := kvdb.collection.Find(ctx, filter)
	if err != nil {
		return nil, err
	}
	return &mongoKVIterator{
		cursor: cursor,
		ctx:    ctx,
	}, nil
}

func (kvdb *mongoKVDB) Close() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	err := kvdb.client.Disconnect(ctx)
	if err != nil {
		gwlog.Errorf("Failed to disconnect from MongoDB: %v", err)
	}
}

func (kvdb *mongoKVDB) IsConnectionError(err error) bool {
	return err == io.EOF || err == io.ErrUnexpectedEOF
}
